!> author: 左志华
!> date: 2022-07-23
!>
!> Mathematics <br>
!> 数学库
!> @note 采用双精度
module easy_math_m

    use, intrinsic :: ieee_arithmetic, only: ieee_is_nan
    use, intrinsic :: iso_fortran_env, only: dp => real64
    implicit none
    private

    public :: arg, argd, cross_product, diff, is_close, Pi, sqrt_eps, &
              arange, linspace

    interface diff
        procedure :: diff_real, diff_integer
    end interface diff

    interface is_close
        procedure :: is_close_real, is_close_cmplx
    end interface is_close

    interface arange
        procedure :: arange_integer, arange_real
    end interface arange

    interface linspace
        module procedure :: linspace_integer, linspace_real
    end interface linspace

    real(dp), parameter :: Pi = acos(-1.0_dp)                   !! Pi <br>
                                                                !! 圆周率
    real(dp), parameter :: sqrt_eps = sqrt(epsilon(1.0_dp))     !! sqrt(epsilon) <br>
                                                                !! 小量，平方根精度
    complex(dp), parameter :: zero_cmplx = (0.0_dp, 0.0_dp)

contains

    !> Calculate cross product <br>
    !> 计算叉积
    pure function cross_product(a, b) result(c)
        real(dp), intent(in) :: a(3), b(3)
        real(dp) :: c(3)

        c(1) = a(2)*b(3) - a(3)*b(2)
        c(2) = a(3)*b(1) - a(1)*b(3)
        c(3) = a(1)*b(2) - a(2)*b(1)

    end function cross_product

    !> calculate argument of complex number <br>
    !> 计算复数的幅角（rad）
    elemental function arg(z) result(a)
        complex(dp), intent(in) :: z
        real(dp) :: a

        if (z == zero_cmplx) then
            a = 0.0_dp
        else
            a = atan2(z%im, z%re)
        end if

    end function arg

    !> calculate argument of complex number <br>
    !> 计算复数的幅角（°）
    elemental function argd(z) result(a)
        complex(dp), intent(in) :: z
        real(dp) :: a

        if (z == zero_cmplx) then
            a = 0.0_dp
        else
            a = atan2(z%im, z%re)*180.0_dp/Pi
        end if

    end function argd

    !> calculate difference of x <br>
    !> 计算x的差值
    pure function diff_real(x, n, prepend, append) result(y)
        real(dp), intent(in) :: x(:)
        integer, intent(in), optional :: n
        real(dp), intent(in), optional :: prepend(:), append(:)
        real(dp), allocatable :: y(:)
        integer :: size_prepend, size_append, size_x, size_work
        integer :: n_, i

        if (present(n)) then
            n_ = n
        else
            n_ = 1
        end if

        if (n_ <= 0) then
            y = x
            return
        end if

        size_prepend = 0
        size_append = 0
        if (present(prepend)) size_prepend = size(prepend)
        if (present(append)) size_append = size(append)
        size_x = size(x)
        size_work = size_x + size_prepend + size_append

        if (size_work <= n_) then
            allocate (y(0))
            return
        end if

        !> Use a quick exit for the common case, to avoid memory allocation.
        if (size_prepend == 0 .and. size_append == 0 .and. n_ == 1) then
            y = x(2:) - x(1:size_x - 1)
            return
        end if

        block
            real(dp) :: work(size_work)
            if (size_prepend > 0) work(:size_prepend) = prepend
            work(size_prepend + 1:size_prepend + size_x) = x
            if (size_append > 0) work(size_prepend + size_x + 1:) = append

            do i = 1, n_
                work(1:size_work - i) = work(2:size_work - i + 1) - work(1:size_work - i)
            end do

            y = work(1:size_work - n_)
        end block

    end function diff_real

    !> calculate difference of x <br>
    !> 计算x的差值
    pure function diff_integer(x, n, prepend, append) result(y)
        integer, intent(in) :: x(:)
        integer, intent(in), optional :: n
        integer, intent(in), optional :: prepend(:), append(:)
        integer, allocatable :: y(:)
        integer :: size_prepend, size_append, size_x, size_work
        integer :: n_, i

        if (present(n)) then
            n_ = n
        else
            n_ = 1
        end if

        if (n_ <= 0) then
            y = x
            return
        end if

        size_prepend = 0
        size_append = 0
        if (present(prepend)) size_prepend = size(prepend)
        if (present(append)) size_append = size(append)
        size_x = size(x)
        size_work = size_x + size_prepend + size_append

        if (size_work <= n_) then
            allocate (y(0))
            return
        end if

        !> Use a quick exit for the common case, to avoid memory allocation.
        if (size_prepend == 0 .and. size_append == 0 .and. n_ == 1) then
            y = x(2:) - x(1:size_x - 1)
            return
        end if

        block
            integer :: work(size_work)
            if (size_prepend > 0) work(:size_prepend) = prepend
            work(size_prepend + 1:size_prepend + size_x) = x
            if (size_append > 0) work(size_prepend + size_x + 1:) = append

            do i = 1, n_
                work(1:size_work - i) = work(2:size_work - i + 1) - work(1:size_work - i)
            end do

            y = work(1:size_work - n_)
        end block

    end function diff_integer

    !> check if a and b are close <br>
    !> 检查a和b是否相近
    elemental logical function is_close_real(a, b, rel_tol, abs_tol, equal_nan) result(close)
        real(dp), intent(in) :: a, b
        real(dp), intent(in), optional :: rel_tol, abs_tol
        logical, intent(in), optional :: equal_nan
        logical :: equal_nan_
        real(dp) :: abs_tol_, rel_tol_

        if (present(equal_nan)) then
            equal_nan_ = equal_nan
        else
            equal_nan_ = .false.
        end if

        if (ieee_is_nan(a) .or. ieee_is_nan(b)) then
            close = merge(.true., .false., equal_nan_ .and. ieee_is_nan(a) .and. ieee_is_nan(b))
        else

            if (present(rel_tol)) then
                rel_tol_ = rel_tol
            else
                rel_tol_ = sqrt_eps
            end if

            if (present(abs_tol)) then
                abs_tol_ = abs_tol
            else
                abs_tol_ = 0.0_dp
            end if

            close = abs(a - b) <= max(abs(rel_tol_*max(abs(a), abs(b))), &
                                      abs(abs_tol_))
        end if

    end function is_close_real

    !> check if a and b are close <br>
    !> 检查 a 和 b 是否相近
    elemental logical function is_close_cmplx(a, b, rel_tol, abs_tol, equal_nan) result(close)
        complex(dp), intent(in) :: a, b
        real(dp), intent(in), optional :: rel_tol, abs_tol
        logical, intent(in), optional :: equal_nan

        close = is_close_real(a%re, b%re, rel_tol, abs_tol, equal_nan) .and. &
                is_close_real(a%im, b%im, rel_tol, abs_tol, equal_nan)

    end function is_close_cmplx

    !> constructs a vector of integers from start to end with step <br>
    !> 从 start 到 end 构造一个整数向量，步长为 step
    pure function arange_integer(start, end, step) result(v)
        integer, intent(in) :: start
        integer, intent(in), optional :: end, step
        integer, allocatable :: v(:)

        integer start_, end_, step_
        integer i

        if (present(end)) then
            start_ = start
            end_ = end
        else
            start_ = 1
            end_ = start
        end if

        if (present(step)) then
            step_ = step
        else
            step_ = 1
        end if

        if (step_ /= 0) then
            step_ = sign(step_, end_ - start_)
        else
            step_ = sign(1, end_ - start_)
        end if

        allocate (v((end_ - start_)/step_ + 1))
        v = [(i, i=start_, end_, step_)]

    end function arange_integer

    !> constructs a vector of reals from start to end with step <br>
    !> 从 start 到 end 构造一个实数向量，步长为 step
    pure function arange_real(start, end, step) result(v)
        real(dp), intent(in) :: start
        real(dp), intent(in), optional :: end, step
        real(dp), allocatable :: v(:)

        real(dp) :: start_, end_, step_
        integer :: i

        if (present(end)) then
            start_ = start
            end_ = end
        else
            start_ = 1.0_dp
            end_ = start
        end if

        if (present(step)) then
            step_ = step
        else
            step_ = 1.0_dp
        end if

        if (step_ /= 0.0_dp) then
            step_ = sign(step_, end_ - start_)
        else
            step_ = sign(1.0_dp, end_ - start_)
        end if

        allocate (v(floor((end_ - start_)/step_) + 1))
        v = [(start_ + (i - 1)*step_, i=1, size(v), 1)]

    end function arange_real

    !> constructs a vector of n linearly spaced numbers from start to end <br>
    !> 从 start 到 end 构造一个 n 个线性间隔的数字向量
    pure function linspace_real(start, end, n) result(v)
        real(dp), intent(in) :: start, end
        integer, intent(in) :: n
        real(dp) :: v(max(n, 0))

        real(dp) :: step
        integer :: i

        if (n <= 0) then
            return
        elseif (n == 1) then
            v(1) = start
            return
        end if

        step = (end - start)/(n - 1)
        do concurrent(i=1:n)
            v(i) = start + (i - 1)*step
        end do

    end function linspace_real

    !> returns a vector of n linearly spaced integers from start to end <br>
    !> 从 start 到 end 构造一个 n 个线性间隔的整数向量
    pure function linspace_integer(start, end, n) result(v)
        integer, intent(in) :: start, end
        integer, intent(in) :: n
        integer :: v(max(n, 0))

        integer :: step
        integer :: i

        if (n <= 0) then
            return
        elseif (n == 1) then
            v(1) = start
            return
        end if

        step = (end - start)/(n - 1)
        do concurrent(i=1:n)
            v(i) = start + (i - 1)*step
        end do

    end function linspace_integer

end module easy_math_m
